import cv2
import numpy as np
import onnxruntime
from scipy.spatial.transform import Rotation as R_scipy
import math
import func


CLASSES = ['jiuwan', 'liuwan', 'siwan', 'wuwan'] 
# sigmod激活函数
def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# 将预测框的格式从中心+框大小改为左上+右下坐标格式
def xywh2xyxy(box):
    x, y, w, h = box.T
    x1 = x - w / 2
    y1 = y - h / 2
    x2 = x + w / 2
    y2 = y + h / 2
    return np.stack([x1, y1, x2, y2], axis=1)

def process_masks(mask_coeffs, protos, shape, threshold=0.5):
    """
    参数：
        - mask_coeffs: [N, 32]
        - protos: [1, 32, 160, 160]
        - shape: 原图尺寸（H, W）
    返回：
        - masks: [N, H, W] 二值图
    """
    protos = protos.squeeze(0)  # [32, 160, 160]
    c, mh, mw = protos.shape

    # 点积系数 × proto → 得到掩码
    masks = sigmoid(np.matmul(mask_coeffs, protos.reshape(c, -1)))  # [N, 25600]
    masks = masks.reshape(-1, mh, mw)  # [N, 160, 160]

    # resize 到原图大小
    masks = np.array([cv2.resize(m, (shape[1], shape[0])) for m in masks])  # shape=(N, H, W)
    masks = masks > threshold  # 转成二值掩码
    return masks.astype(np.uint8)

def apply_nmc(boxes, scores, score_threshold=0.5, iou_threshold=0.5):
    if len(boxes) == 0:
        return []
    boxes_xywh = np.array(boxes).astype(np.float32).tolist()
    scores = scores.astype(np.float32).tolist()

    # NMSBoxes 返回的是 [[i], [j], ...] 嵌套形式
    indices = cv2.dnn.NMSBoxes(boxes_xywh, scores, score_threshold, iou_threshold)

    # 展平
    if len(indices) > 0:
        indices = indices.flatten()
    else:
        indices = []

    return indices

def crop_mask(masks, boxes):
    """
    NumPy 实现：将每个掩码裁剪到对应的框内部，其它区域置为0。
    
    参数：
        masks: np.ndarray，形状为 (n, h, w)，每个掩码
        boxes: np.ndarray，形状为 (n, 4)，格式为 [x1, y1, x2, y2]，坐标是浮点型（像素）
    
    返回：
        裁剪后的 masks，形状为 (n, h, w)
    """
    n, h, w = masks.shape
    out = np.zeros_like(masks)

    for i in range(n):
        x1, y1, x2, y2 = boxes[i]
        # 转为整数像素坐标，并裁剪范围
        x1 = max(int(np.floor(x1)), 0)
        y1 = max(int(np.floor(y1)), 0)
        x2 = min(int(np.ceil(x2)), w)
        y2 = min(int(np.ceil(y2)), h)

        out[i, y1:y2, x1:x2] = masks[i, y1:y2, x1:x2]

    return out

def postprocess(output0, output1, orig_shape, conf_thres=0.8):
    """
    参数：
        - output0: 模型主输出 [1, 25200, 38]
        - output1: 掩码原型输出 [1, 32, 160, 160]
        - orig_shape: 原图尺寸（H, W）
    返回：
        - boxes: [N, 4] xyxy
        - scores: [N]
        - classes: [N]
        - masks: [N, H, W]
    """
    pred = output0.squeeze(0)  # [25200, 38]
    boxes = pred[:, :4]              # [cx, cy, w, h]
    objectness = pred[:, 4]          # 目标置信度
    class_score = pred[:, 5]         # 只有1类
    mask_coeffs = pred[:, 6:]        # [25200, 32]

    # 置信度 = objectness × class_score
    scores = objectness * class_score
    keep0 = scores > conf_thres

    # NMS
    ind2 = apply_nmc(boxes, scores)

    ind1 = np.where(keep0)[0]
    keep = np.intersect1d(ind1, ind2)
    try:
        boxes = boxes[keep]
        scores = scores[keep]
        classes = np.zeros_like(scores)  # 只有1类，直接赋值为0
        mask_coeffs = mask_coeffs[keep]

        # 转换为 xyxy 坐标
        boxes = xywh2xyxy(boxes)

        # 还原掩码
        masks = process_masks(mask_coeffs, output1, orig_shape)

        masks = crop_mask(masks, boxes)
    except:
        boxes, scores, classes, masks = []

    return boxes, scores, classes, masks

def postprocess1(output, conf_thres=0.5, nms_thres=0.5):
    """
    参数：
        - output: 模型输出 [1, 25200, 9]
        - orig_shape: 原图大小 (H, W)
    返回：
        - boxes: [N, 4] in xyxy format
        - scores: [N]
        - classes: [N]
    """
    pred = output.squeeze(0)  # [25200, 9]
    boxes = pred[:, :4]       # [cx, cy, w, h]
    obj_conf = pred[:, 4:5]   # [25200, 1]
    class_confs = pred[:, 5:] # [25200, num_classes]

    # 每个 anchor 选择置信度最高的类别及得分
    class_scores = np.max(class_confs, axis=1)     # [25200]
    class_ids = np.argmax(class_confs, axis=1)     # [25200]
    scores = (obj_conf[:, 0]) * class_scores       # [25200]

    # 置信度阈值过滤
    keep = scores > conf_thres
    boxes = boxes[keep]
    scores = scores[keep]
    class_ids = class_ids[keep]

    # 坐标转换
    boxes = xywh2xyxy(boxes)

    # NMS
    keep_indices = apply_nmc(boxes, scores)

    return boxes[keep_indices], scores[keep_indices], class_ids[keep_indices]


def mask_to_bgr_mask(mask, color=(0, 255, 0), alpha=0.5):
    """
    将单通道 mask 转为三通道 BGR 彩色图像，用于可视化。
    参数：
        mask: (H, W) numpy array，0~1 或 0/1 掩码
        color: (B, G, R)，你想给掩码涂的颜色
        alpha: 颜色强度系数（用于叠加）
    返回：
        bgr_mask: (H, W, 3) 彩色掩码图像，uint8
    """
    mask = mask.astype(np.float32)
    mask = np.clip(mask, 0.0, 1.0)

    bgr_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
    for i in range(3):
        bgr_mask[:, :, i] = (mask * color[i] * alpha).astype(np.uint8)

    return bgr_mask

def findcorner(img, frame):
    # 寻找边缘
    contours, a = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    contour = max(contours, key=cv2.contourArea)

    # 下面是画轮廓的代码
    # image = cv2.drawContours(frame, contour, -1, (0, 255, 0), 1)
    # cv2.imshow('img', image)
    # cv2.waitKey(0)

    # 对正方体进行逼近多边形
    dots = cv2.approxPolyDP(contour, 12, True)
    dots = dots.reshape(-1, 2)


    if len(dots) != 4:
        rect = cv2.minAreaRect(contour)        # 返回 (中心坐标(x,y), (宽,高), 角度)
        box = cv2.boxPoints(rect)              
        box = box.astype(np.uint32)
        boxed = box[[1, 0, 3, 2]]
        return boxed
    else:
        return dots

def get_position(dots, param):
    # 求解pnp问题
    pts_2d = np.array(dots, dtype=np.float32)
    K, dist_coeffs = np.load("720p.npz").values()
    l, w, h = (46, 34, 12)
    pts_3d = np.array([
        [-w / 2, l / 2, h / 2],
        [w / 2, l / 2, h / 2],
        [w / 2, -l / 2, h / 2],
        [-w / 2, -l / 2, h / 2]
    ], dtype=np.float32)*param
    success, rvec, tvec = cv2.solvePnP(pts_3d, pts_2d, K, dist_coeffs)
    # 1. 转换 rvec → 旋转矩阵 R (3×3)
    R, _ = cv2.Rodrigues(rvec)

    # 2. 拼接成 3×4 矩阵
    Rt = np.hstack((R, tvec))  # [R | t]

    # 3. 扩展为 4×4 齐次矩阵
    T_ = np.vstack((Rt, np.array([[0, 0, 0, 1]])))
    return T_

def trans_to_arm(T1, T2):
    T_ = np.multiply(T1, T2)
    R = T_[:3, :3]
    t = T_[:3, 3]
    rot = R_scipy.from_matrix(R)
    euler = rot.as_euler('xyz', degrees=True)  # 或 'zyx', 'yxz' 等
    print("欧拉角（角度制）:", euler)
    return euler, t

def euler_to_trans(euler_angles, t):
    r = R_scipy.from_euler('xyz', euler_angles, degrees=True)
    R_matrix = r.as_matrix()
    t = t.reshape(3, 1)
    T_ = np.vstack((np.hstack((R_matrix, t)), np.array([[0, 0, 0, 1]])))
    return T_

# 未实现的方法二
def hough_line_to_params(rho, theta):
    a = np.cos(theta)
    b = np.sin(theta)
    x0 = a * rho
    y0 = b * rho

    # 两个点
    x1 = int(x0 + 1000 * (-b))
    y1 = int(y0 + 1000 * (a))
    x2 = int(x0 - 1000 * (-b))
    y2 = int(y0 - 1000 * (a))

    # 转为直线标准形式 Ax + By + C = 0
    A = y2 - y1
    B = x1 - x2
    C = x2 * y1 - x1 * y2
    return A, B, -C, (x1, y1), (x2, y2)

def hough_lines(img):
    # img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(img, 50, 150)
    lines = cv2.HoughLines(edges, 1, np.pi / 180, threshold=100)
    for l in lines:
        rho, theta = l[0]
        # params = hough_line_to_params(rho, theta)
        a, b, c, d1, d2 = hough_line_to_params(rho, theta)
    cv2.line(img, d1, d2, (155, 155, 155), 2)


CAMERA_POSITION = np.array([[0, 0, math.pi/4], [-2, 0, -5]], dtype=np.float64)



# 绘制box
# for box in boxes:
#     x1 = int(box[0])
#     y1 = int(box[1])
#     x2 = int(box[2])
#     y2 = int(box[3])
#     cv2.rectangle(orig_img, (x1, y1), (x2, y2), (0, 255, 0), 2)

# 掩膜绘制
# for mask in masks:
#     msk = mask_to_bgr_mask(mask)
#     result = cv2.addWeighted(orig_img, 1.0, msk, 1.0, 0)
#     cv2.imshow("1", result)
#     cv2.waitKey()
# cv2.destroyAllWindows()

def run(input_img, path, show=False, mode="segment"):
    # 定义模型并给出输出
    model = func.YOLOV5(path)
    frame = input_img
    outputs, orig_img = model.inference(frame)
    if mode == "segment":
        output0, output1 = outputs
        boxes, scores, classes, masks = postprocess(output0, output1, orig_img.shape[:2])
        majiang = masks[0]
        positions = []
        for majiang in masks:
            majiang = majiang*255
            majiang_ori = func.restore_image(majiang)
            _, majiang_ori = cv2.threshold(majiang_ori, 10, 255, cv2.THRESH_BINARY)
            # hough_lines(majiang_ori)

            dots = findcorner(majiang_ori, frame)
            positions.append(dots)
        positions = np.array(positions, dtype=np.float32)
        # 提取所有 x 坐标
        x_values = positions[:, :, 0]  # shape: (n, 4)
        # 找到全局最小 x 的位置
        min_quad_idx, _ = np.unravel_index(np.argmin(x_values), x_values.shape)
        # 取出该角点所在的整个四边形
        min_quad = positions[min_quad_idx]

        if show:
            # try:
            for dot in min_quad:
                cv2.circle(frame, tuple(dot), 1, (0, 255, 0), 2)
            # except:
            #     pass
            cv2.imshow("", frame)
            cv2.waitKey(1)
        
    elif mode == "detect":
        boxes, scores, classes  = postprocess1(outputs)

        if show:
            try:
                # 绘制box
                for box, score, cl in zip(boxes, scores, classes):   
                    x1 = int(box[0])
                    y1 = int(box[1])
                    x2 = int(box[2])
                    y2 = int(box[3])
                    cv2.rectangle(orig_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(orig_img, '{0} {1:.2f}'.format(CLASSES[cl], score), (x1, y1),
                                cv2.FONT_HERSHEY_SIMPLEX,
                                0.6, (0, 0, 255), 2)
                    frame = func.restore_image(orig_img)
            except:
                print("no object")
            cv2.imshow("0", frame)
            cv2.waitKey(1)

    # M1 = get_position(dots, 1)
    # M2 = euler_to_trans(*CAMERA_POSITION)
    # pose = trans_to_arm(M1, M2)
    # return pose

if __name__ == "__main__":
    img = cv2.imread("onnxdetect/111.jpg")
    run(img, "onnxdetect/bestseg.onnx", True, "segment")
    cv2.waitKey()
    cv2.destroyAllWindows()